{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf7ccb32",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/langchain-ai/langchain-academy/blob/main/module-2/chatbot-external-memory.ipynb) [![Open in LangChain Academy](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66e9eba12c7b7688aa3dbb5e_LCA-badge-green.svg)](https://academy.langchain.com/courses/take/intro-to-langgraph/lessons/58239440-lesson-6-chatbot-w-summarizing-messages-and-external-memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af6c7afe-1037-41ab-98e4-494692e47402",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Chatbot with message summarization & external DB memory\n",
    "\n",
    "## Review\n",
    "\n",
    "We've covered how to customize graph state schema and reducer. \n",
    " \n",
    "We've also shown a number of tricks for trimming or filtering messages in graph state. \n",
    "\n",
    "We've used these concepts in a Chatbot with memory that produces a running summary of the conversation.\n",
    "\n",
    "## Goals\n",
    "\n",
    "But, what if we want our Chatbot to have memory that persists indefinitely?\n",
    "\n",
    "Now, we'll introduce some more advanced checkpointers that support external databases. \n",
    "\n",
    "Here, we'll show how to use [Sqlite as a checkpointer](https://docs.langchain.com/oss/python/langgraph/persistence#checkpointer-libraries), but other checkpointers, such as  Postgres are available!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85ed78d9-6ca2-45ac-96a9-52e341ec519d",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph-checkpoint-sqlite langchain_core langgraph langchain_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e10c4d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, getpass\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b40d25c0-e9b5-4854-bf07-3cc3ff07122e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Sqlite\n",
    "\n",
    "A good starting point here is the [SqliteSaver checkpointer](https://docs.langchain.com/oss/python/langgraph/persistence#checkpointer-libraries).\n",
    "\n",
    "Sqlite is a [small, fast, highly popular](https://x.com/karpathy/status/1819490455664685297) SQL database. \n",
    " \n",
    "If we supply `\":memory:\"` it creates an in-memory Sqlite database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fae15402-17ae-4e89-8ecf-4c89e08b22fe",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "# In memory\n",
    "conn = sqlite3.connect(\":memory:\", check_same_thread = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2bf53ec-6d4a-42ce-8183-344795eed403",
   "metadata": {},
   "source": [
    "But, if we supply a db path, then it will create a database for us!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58339167-920c-4994-a0a7-0a9c5d4f7cf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pull file if it doesn't exist and connect to local db\n",
    "!mkdir -p state_db && [ ! -f state_db/example.db ] && wget -P state_db https://github.com/langchain-ai/langchain-academy/raw/main/module-2/state_db/example.db\n",
    "\n",
    "db_path = \"state_db/example.db\"\n",
    "conn = sqlite3.connect(db_path, check_same_thread=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7736b6-a750-48f8-a838-8e7616b12250",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Here is our checkpointer \n",
    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
    "memory = SqliteSaver(conn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d8cb629-213f-4b87-965e-19b812c42da1",
   "metadata": {},
   "source": [
    "Let's re-define our chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc414e29-2078-41a0-887c-af1a6a3d72c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import Literal\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import SystemMessage, HumanMessage, RemoveMessage\n",
    "\n",
    "from langgraph.graph import END\n",
    "from langgraph.graph import MessagesState\n",
    "\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\",temperature=0)\n",
    "\n",
    "class State(MessagesState):\n",
    "    summary: str\n",
    "\n",
    "# Define the logic to call the model\n",
    "def call_model(state: State):\n",
    "    \n",
    "    # Get summary if it exists\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "\n",
    "    # If there is summary, then we add it\n",
    "    if summary:\n",
    "        \n",
    "        # Add summary to system message\n",
    "        system_message = f\"Summary of conversation earlier: {summary}\"\n",
    "\n",
    "        # Append summary to any newer messages\n",
    "        messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
    "    \n",
    "    else:\n",
    "        messages = state[\"messages\"]\n",
    "    \n",
    "    response = model.invoke(messages)\n",
    "    return {\"messages\": response}\n",
    "\n",
    "def summarize_conversation(state: State):\n",
    "    \n",
    "    # First, we get any existing summary\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "\n",
    "    # Create our summarization prompt \n",
    "    if summary:\n",
    "        \n",
    "        # A summary already exists\n",
    "        summary_message = (\n",
    "            f\"This is summary of the conversation to date: {summary}\\n\\n\"\n",
    "            \"Extend the summary by taking into account the new messages above:\"\n",
    "        )\n",
    "        \n",
    "    else:\n",
    "        summary_message = \"Create a summary of the conversation above:\"\n",
    "\n",
    "    # Add prompt to our history\n",
    "    messages = state[\"messages\"] + [HumanMessage(content=summary_message)]\n",
    "    response = model.invoke(messages)\n",
    "    \n",
    "    # Delete all but the 2 most recent messages\n",
    "    delete_messages = [RemoveMessage(id=m.id) for m in state[\"messages\"][:-2]]\n",
    "    return {\"summary\": response.content, \"messages\": delete_messages}\n",
    "\n",
    "# Determine whether to end or summarize the conversation\n",
    "def should_continue(state: State)-> Literal [\"summarize_conversation\",END]:\n",
    "    \n",
    "    \"\"\"Return the next node to execute.\"\"\"\n",
    "    \n",
    "    messages = state[\"messages\"]\n",
    "    \n",
    "    # If there are more than six messages, then we summarize the conversation\n",
    "    if len(messages) > 6:\n",
    "        return \"summarize_conversation\"\n",
    "    \n",
    "    # Otherwise we can just end\n",
    "    return END"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c13c0b-a383-4f73-9cc1-63f0eed8f190",
   "metadata": {},
   "source": [
    "Now, we just re-compile with our sqlite checkpointer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e867fd95-91eb-4ce1-82fc-bb72d611a96d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langgraph.graph import StateGraph, START\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(State)\n",
    "workflow.add_node(\"conversation\", call_model)\n",
    "workflow.add_node(summarize_conversation)\n",
    "\n",
    "# Set the entrypoint as conversation\n",
    "workflow.add_edge(START, \"conversation\")\n",
    "workflow.add_conditional_edges(\"conversation\", should_continue)\n",
    "workflow.add_edge(\"summarize_conversation\", END)\n",
    "\n",
    "# Compile\n",
    "graph = workflow.compile(checkpointer=memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8769db99-3938-45e6-a594-56beb18d6c45",
   "metadata": {},
   "source": [
    "Now, we can invoke the graph several times. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4094a0-d240-4be8-903a-7d9f605bdc5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a thread\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "# Start conversation\n",
    "input_message = HumanMessage(content=\"hi! I'm Lance\")\n",
    "output = graph.invoke({\"messages\": [input_message]}, config) \n",
    "for m in output['messages'][-1:]:\n",
    "    m.pretty_print()\n",
    "\n",
    "input_message = HumanMessage(content=\"what's my name?\")\n",
    "output = graph.invoke({\"messages\": [input_message]}, config) \n",
    "for m in output['messages'][-1:]:\n",
    "    m.pretty_print()\n",
    "\n",
    "input_message = HumanMessage(content=\"i like the 49ers!\")\n",
    "output = graph.invoke({\"messages\": [input_message]}, config) \n",
    "for m in output['messages'][-1:]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0f3e842-4497-45e2-a924-69672a9bcb33",
   "metadata": {},
   "source": [
    "Let's confirm that our state is saved locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ab158a-5a82-417a-8841-730a4cc18ea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "graph_state = graph.get_state(config)\n",
    "graph_state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e21152d-ed9c-408d-b7d5-f634c9ce81e2",
   "metadata": {},
   "source": [
    "### Persisting state\n",
    "\n",
    "Using database like Sqlite means state is persisted! \n",
    "\n",
    "For example, we can re-start the notebook kernel and see that we can still load from Sqlite DB on disk.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9a44dc5-be04-45fa-a6fc-27b0f8ee4678",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a thread\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "graph_state = graph.get_state(config)\n",
    "graph_state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8466e418-1a46-4cdb-a51a-6ae14281bb85",
   "metadata": {},
   "source": [
    "## Studio\n",
    "\n",
    "**⚠️ Notice**\n",
    "\n",
    "Since filming these videos, we've updated Studio so that it can now be run locally and accessed through your browser. This is the preferred way to run Studio instead of using the Desktop App shown in the video. It is now called _LangSmith Studio_ instead of _LangGraph Studio_. Detailed setup instructions are available in the \"Getting Setup\" guide at the start of the course. You can find a description of Studio [here](https://docs.langchain.com/langsmith/studio), and specific details for local deployment [here](https://docs.langchain.com/langsmith/quick-start-studio#local-development-server).  \n",
    "To start the local development server, run the following command in your terminal in the `/studio` directory in this module:\n",
    "\n",
    "```\n",
    "langgraph dev\n",
    "```\n",
    "\n",
    "You should see the following output:\n",
    "```\n",
    "- 🚀 API: http://127.0.0.1:2024\n",
    "- 🎨 Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024\n",
    "- 📚 API Docs: http://127.0.0.1:2024/docs\n",
    "```\n",
    "\n",
    "Open your browser and navigate to the **Studio UI** URL shown above.\n",
    "Load the `chatbot` in Studio, which uses `module-2/studio/chatbot.py` set in `module-2/studio/langgraph.json`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4916d8b",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b41d4491-612e-4698-aa60-de76277422bb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
